package main

import (
	"fmt"
	"net/http"
	"time"
)

// 长连接处理函数，持续向客户端发送消息
func streamHandler(w http.ResponseWriter, r *http.Request) {
	// 设置响应头，告知客户端这是一个流
	w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("Access-Control-Allow-Origin", "*")

	// 确保客户端不关闭连接
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "Streaming not supported", http.StatusInternalServerError)
		return
	}

	// 客户端连接标识
	clientID := r.URL.Query().Get("clientId")
	if clientID == "" {
		clientID = fmt.Sprintf("client-%d", time.Now().UnixNano())
	}

	fmt.Printf("客户端 %s 已连接\n", clientID)
	defer fmt.Printf("客户端 %s 已断开\n", clientID)

	// 创建一个通道用于监听客户端断开连接
	done := make(chan struct{})
	defer close(done)

	// 监听客户端连接关闭
	go func() {
		<-r.Context().Done()
		done <- struct{}{}
	}()

	// 持续发送消息
	counter := 0
	for {
		select {
		case <-done:
			return
		default:
			// 构建消息
			message := fmt.Sprintf("服务器消息 #%d - %s\n", counter, time.Now().Format(time.RFC3339))

			// 发送消息
			fmt.Fprintf(w, "data: %s\n\n", message)

			// 立即刷新缓冲区，确保客户端能及时收到
			flusher.Flush()

			counter++

			// 每2秒发送一条消息
			time.Sleep(2 * time.Second)
		}
	}
}

// 用于发送自定义消息的处理函数
func sendMessageHandler(w http.ResponseWriter, r *http.Request) {
	// 在实际应用中，这里可以将消息存储到某个队列
	// 然后由streamHandler读取并推送给客户端
	message := r.URL.Query().Get("message")
	if message == "" {
		http.Error(w, "消息不能为空", http.StatusBadRequest)
		return
	}

	fmt.Printf("收到自定义消息: %s\n", message)
	fmt.Fprintf(w, "消息已接收: %s", message)
}

func main() {
	// 注册处理函数
	http.HandleFunc("/stream", streamHandler)
	http.HandleFunc("/send", sendMessageHandler)

	// 启动服务器
	fmt.Println("服务器启动在 :18080 端口")
	if err := http.ListenAndServe(":18080", nil); err != nil {
		fmt.Printf("服务器启动失败: %s\n", err)
	}
}
